package com.yyx.demo.queue;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.ACL;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
import java.util.TreeMap;
import java.util.concurrent.CountDownLatch;

public class ZKDistributedQueue {

    private static final Logger LOG = LoggerFactory.getLogger(ZKDistributedQueue.class);

    private final String dir;

    private ZooKeeper zookeeper;
    private List<ACL> acl = ZooDefs.Ids.OPEN_ACL_UNSAFE;

    private final String prefix = "qn-";

//    private String topic;
//
//    //设置topic
//    public Boolean setTopic(String topic) throws KeeperException, InterruptedException {
//        //判断有无topic 路径
//        Stat stat=zookeeper.exists(topic,null);
//        if(stat==null){
//            zookeeper.create(dir + "/" + prefix, new byte[0], acl, CreateMode.PERSISTENT_SEQUENTIAL);
//            return true;
//        }else {
//            return false;
//        }
//    }

    public ZKDistributedQueue(ZooKeeper zookeeper, String dir, List<ACL> acl) {
        this.dir = dir;
        if (acl != null) {
            this.acl = acl;
        }
        this.zookeeper = zookeeper;
    }

    public byte[] remove() throws NoSuchElementException, KeeperException, InterruptedException {
        Map<Long, String> orderedChildren;
        // Same as for element.  Should refactor this.
        while (true) {
            try {
                orderedChildren = orderedChildren(null);
            } catch (KeeperException.NoNodeException e) {
                throw new NoSuchElementException();
            }
            if (orderedChildren.size() == 0) {
                throw new NoSuchElementException();
            }

            for (String headNode : orderedChildren.values()) {
                String path = dir + "/" + headNode;
                try {
                    byte[] data = zookeeper.getData(path, false, null);
                    zookeeper.delete(path, -1);
                    return data;
                } catch (KeeperException.NoNodeException e) {
                    // 其他客户端先删除了这个节点
                }
            }

        }
    }

    private static class LatchChildWatcher implements Watcher {

        CountDownLatch latch;

        public LatchChildWatcher() {
            latch = new CountDownLatch(1);
        }

        @Override
        public void process(WatchedEvent event) {
            LOG.debug("Watcher fired: {}", event);
            latch.countDown();
        }

        public void await() throws InterruptedException {
            latch.await();
        }

    }


    //删除队头，返回该元素。同步阻塞
    public byte[] takeSync() throws KeeperException, InterruptedException {
        Map<Long, String> orderedChildren;
        // Same as for element.  Should refactor this.
        while (true) {
            LatchChildWatcher childWatcher = new LatchChildWatcher();
            try {
                orderedChildren = orderedChildren(childWatcher);
            } catch (KeeperException.NoNodeException e) {
                zookeeper.create(dir, new byte[0], acl, CreateMode.PERSISTENT);
                continue;
            }
            if (orderedChildren.size() == 0) {
                //拿不到就阻塞
                childWatcher.await();
                continue;
            }
            for (String headNode : orderedChildren.values()) {
                String path = dir + "/" + headNode;
                try {
                    byte[] data = zookeeper.getData(path, false, null);
                    zookeeper.delete(path, -1);
                    return data;
                } catch (KeeperException.NoNodeException e) {
                    // Another client deleted the node first.
                }
            }
        }
    }


    //异步读
    public void  takeAsync(TakeCommitCallback callBack) throws KeeperException, InterruptedException {
        Map<Long, String> orderedChildren;
        // Same as for element.  Should refactor this.
        while (true) {
            LatchChildWatcher childWatcher = new LatchChildWatcher();
            try {
                orderedChildren = orderedChildren(childWatcher);
            } catch (KeeperException.NoNodeException e) {
                zookeeper.create(dir, new byte[0], acl, CreateMode.PERSISTENT);
                continue;
            }
            if (orderedChildren.size() == 0) {
                //拿不到就阻塞
                childWatcher.await();
                continue;
            }

            for (String headNode : orderedChildren.values()) {
                String path = dir + "/" + headNode;
                try {
                    byte[] data = zookeeper.getData(path, false, null);
                    zookeeper.delete(path, -1);
                    callBack.takeComplete(data);
                } catch (KeeperException.NoNodeException e) {
                    // Another client deleted the node first.
                }
            }
        }
    }




    //以id构成的结构体，存元
    private Map<Long, String> orderedChildren(Watcher watcher) throws KeeperException, InterruptedException {
        Map<Long, String> orderedChildren = new TreeMap<>();

        List<String> childNames;
        childNames = zookeeper.getChildren(dir, watcher);

        for (String childName : childNames) {
            try {
                //检查名字是否一样,避免错误ZK操作
                if (!childName.regionMatches(0, prefix, 0, prefix.length())) {
                    LOG.warn("Found child node with improper name: {}", childName);
                    continue;
                }
                String suffix = childName.substring(prefix.length());
                Long childId = Long.parseLong(suffix);
                orderedChildren.put(childId, childName);
            } catch (NumberFormatException e) {
                LOG.warn("Found child node with improper format : {}", childName, e);
            }
        }

        return orderedChildren;
    }



    //同步提交
    public boolean offerSync(byte[] data) throws KeeperException, InterruptedException {
        while (true) {
            try {
                zookeeper.create(dir + "/" + prefix, data, acl, CreateMode.PERSISTENT_SEQUENTIAL);
                return true;
                //节点不存在则创建节点
            } catch (KeeperException.NoNodeException e) {
                zookeeper.create(dir, new byte[0], acl, CreateMode.PERSISTENT);
            }
        }
    }

    //异步提交
    public void offerAync(byte[] data,OfferCommitCallback offerCommitCallback) throws KeeperException, InterruptedException {
        while (true) {
            try {
                zookeeper.create(dir + "/" + prefix, data, acl, CreateMode.PERSISTENT_SEQUENTIAL);
                offerCommitCallback.offerComplete(data);
                break;
                //节点不存在则创建节点
            } catch (KeeperException.NoNodeException e) {
                zookeeper.create(dir, new byte[0], acl, CreateMode.PERSISTENT);
            }
        }
    }
}
